package com.qzdatasoft.ratelimit.core;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Map;
import java.util.Queue;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

public class LeakyBucketLimiter {
    private static final Logger logger = LoggerFactory.getLogger(LeakyBucketLimiter.class);
    private final CustomRateLimitRule props; // 自定义限流配置
    private final Queue<RequestTask> queue = new ConcurrentLinkedQueue<>(); // 请求队列
    private final AtomicInteger centralCount = new AtomicInteger(0); // 总计数器
    private final Integer compensationNumber = 0; // 补偿数量
//    private final Map<String, AtomicInteger> urlCounts = new ConcurrentHashMap<>();
//    private final Map<String, AtomicInteger> ipCounts = new ConcurrentHashMap<>();
//    private final Map<String, AtomicInteger> userCounts = new ConcurrentHashMap<>();

    private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor(r -> {
        Thread t = new Thread(r, "leaky-bucket-consumer");
        t.setDaemon(true);
        return t;
    });

    public CustomRateLimitRule getProps() {
        return props;
    }

    public LeakyBucketLimiter(CustomRateLimitRule props) {
        this.props = props;
        long interval = props.getIntervalNanos();
        scheduler.scheduleAtFixedRate(this::drainOne, interval, interval, TimeUnit.NANOSECONDS);

        logger.info("{} initialized: {}", this.getClass().getSimpleName(), props.toString());
    }

    public boolean tryEnqueue(RequestTask task, String resourceKey/*, String url, String ip, String user*/) {
        // check capacities
        if (centralCount.get() >= props.getCapacity()) return false;// 超容
//        if (getCount(urlCounts, url) >= props.getPerUrlCapacity()) return false;
//        if (getCount(ipCounts, ip) >= props.getPerIpCapacity()) return false;
//        if (getCount(userCounts, user) >= props.getPerUserCapacity()) return false;

        // increment counts atomically
        centralCount.incrementAndGet();
//        incCount(urlCounts, url);
//        incCount(ipCounts, ip);
//        incCount(userCounts, user);

//        task.setKeys(url, ip, user);
        task.setResourceKey(resourceKey);
        queue.add(task);
        return true;
    }

    private int getCount(Map<String, AtomicInteger> map, String key) {
        AtomicInteger ai = map.get(key);
        return ai == null ? 0 : ai.get();
    }
    private void incCount(Map<String, AtomicInteger> map, String key) {
        map.compute(key, (k, v) -> {
            if (v == null) return new AtomicInteger(1);
            v.incrementAndGet();
            return v;
        });
    }
    private void decCount(Map<String, AtomicInteger> map, String key) {
        map.computeIfPresent(key, (k, v) -> {
            if (v.decrementAndGet() <= 0) return null;
            return v;
        });
    }

    // 上次消费时间
    private volatile long lastDrainTime = System.nanoTime();// 上次消费时间

    // 消费队列中的一个请求任务，
    //
    // 具体工作原理：从队列中取出一个请求任务进行处理。
    // - 如果任务为空则直接返回；
    // - 否则检查该任务在队列中等待的时间是否超时，若超时则拒绝并释放计数；未超时则允许执行并释放计数。
    // 整个过程捕获异常以保证调度器持续运行。
    private void drainOne() {
        try {
            long now = System.nanoTime();
            long timePassed = now - lastDrainTime;

            // 计算自上次滴水以来应处理的滴水数
            long expectedTasks = timePassed / props.getIntervalNanos();

            // 更新最后处理时间
            if(expectedTasks == 0) {
                expectedTasks = 1;
                lastDrainTime = now;
            } else {
                long remainingTime = timePassed % props.getIntervalNanos();
                lastDrainTime = now - remainingTime;
            }

            for (int i = 0; i < expectedTasks && !queue.isEmpty(); i++) {
                RequestTask task = queue.poll();
                if (task == null) continue;

                if (now - task.getEnqueueTime() > props.getTimeoutNanos()) {
                    // 超时处理
                    releaseCounts(task);
                    task.tryCompleteReject();
                } else {
                    // 正常处理
                    releaseCounts(task);
                    task.tryCompleteAllow();
                }
            }
        } catch (Throwable t) {
            logger.warn("Error in drainOne operation", t);
            // TODO: 根据异常类型决定是否继续运行

        }
    }

    // 当等待线程超时并主动移除任务时调用，或 consumer 已处理过则保证幂等
    public boolean removeTask(RequestTask task) {
        boolean removed = queue.remove(task);
        if (removed) {
            releaseCounts(task);
            task.tryCompleteReject();
        }
        return removed;
    }

    private void releaseCounts(RequestTask task) {
        // decrement counts once per task; careful with races: counts were incremented on enqueue
        if (task == null) return;
        centralCount.decrementAndGet();
//        if (task.getUrl() != null) decCount(urlCounts, task.getUrl());
//        if (task.getIp() != null) decCount(ipCounts, task.getIp());
//        if (task.getUser() != null) decCount(userCounts, task.getUser());
    }

    public void shutdown() {
        scheduler.shutdownNow();
    }

}